# check_module.py
# usage:
# python ./src/tools/check_module.py --module tmp_question2sql.json --samples "./data/samples.merge.jsonl"

import argparse
import json
from typing import List, Dict, Set
from clean_question import clean_question

def load_module_questions(module_file: str) -> Set[str]:
    """
    从 --module 文件中加载所有 questions，并返回一个集合（已去重）
    """
    questions = set()
    with open(module_file, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                data = json.loads(line)
                if 'questions' in data:
                    for q in data['questions']:
                        if q.strip():
                            questions.add(q.strip())
            except json.JSONDecodeError as e:
                print(f"Warning: JSON decode error at line {line_num} in module file: {e}")
    return questions

def process_samples(samples_file: str, module_questions: Set[str]):
    """
    顺序读取 samples 文件，clean 每个 question，
    并标记 module_questions 中是否出现过。
    最终输出未被标记（即未在 samples 中出现）的原始问题。
    """
    matched_questions = set()

    with open(samples_file, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                data = json.loads(line)
                raw_question = data.get('question', '')
                if not raw_question:
                    continue

                # 调用 clean_question 清理
                cleaned_sample_q = clean_question(raw_question)

                # 检查 cleaned_sample_q 是否匹配 module 中的任意原始问题（清理后）
                for orig_q in module_questions:
                    if clean_question(orig_q) == cleaned_sample_q:
                        matched_questions.add(orig_q)
                        break  # 匹配到一个即可

            except json.JSONDecodeError as e:
                print(f"Warning: JSON decode error at line {line_num} in samples file: {e}")

    # 找出未被匹配的问题（原始字符串）
    unmatched = module_questions - matched_questions

    # 输出未匹配的问题
    if unmatched:
        print("未被匹配到的问题（来自 --module 的 questions）:")
        for q in sorted(unmatched):  # 排序便于查看
            print(q)
    else:
        print("所有问题均已被匹配。")

def main():
    parser = argparse.ArgumentParser(description="检查 module 中的问题是否在 samples 中出现过（基于 cleaned question）")
    parser.add_argument('--module', required=True, help='module 文件路径（JSONL 格式）')
    parser.add_argument('--samples', required=True, help='samples 文件路径（JSONL 格式）')

    args = parser.parse_args()

    # 加载 module 中的所有原始问题
    print(f"Loading module questions from: {args.module}")
    module_questions = load_module_questions(args.module)
    print(f"Loaded {len(module_questions)} unique questions from module.")

    # 处理 samples 并标记匹配项
    print(f"Processing samples from: {args.samples}")
    process_samples(args.samples, module_questions)

if __name__ == '__main__':
    main()